package org.nowcoder.leetcode.sort;

import java.util.Comparator;
import java.util.Map;
import java.util.HashMap;
import java.util.PriorityQueue;

/**
 * Title  : 347. Top K Frequent Elements
 * Source : https://leetcode.cn/problems/top-k-frequent-elements/description
 * Author : XrazYang
 * Date   : 2023-09-20
 */


public class LeetCode_347 {
    private class TopK {
        private int ele;
        private int count;

        protected TopK(int ele, int count) {
            this.ele = ele;
            this.count = count;
        }
    }

    public int[] topKFrequent(int[] nums, int k) {
        Map<Integer, Integer> map = new HashMap<>();
        //统计频率
        for (int num : nums) {
            Integer count = map.getOrDefault(num, 0);
            map.put(num, count + 1);
        }
        //构建最小堆
        PriorityQueue<TopK> minHeap = new PriorityQueue<>(k, new Comparator<TopK>() {
            @Override
            public int compare(TopK o1, TopK o2) {
                return o1.count - o2.count;
            }
        });
        //Top K
        for (Map.Entry<Integer, Integer> entry : map.entrySet()) {
            TopK topK = new TopK(entry.getKey(), entry.getValue());
            if (minHeap.size() < k) {
                minHeap.offer(topK);
            } else {
                TopK top = minHeap.peek();
                if (top.count < topK.count) {
                    minHeap.poll();
                    minHeap.offer(topK);
                }
            }
        }
        int i = 0;
        int[] elements = new int[k];
        for (TopK ele : minHeap) {
            elements[i++] = ele.ele;
        }
        return elements;
    }
}
